package com.wb.utils

import com.fasterxml.jackson.core.type.TypeReference
import com.fasterxml.jackson.databind.ObjectMapper
import com.google.gson.Gson
import com.google.gson.reflect.TypeToken
import com.wb.component.Slf4j.Companion.log
import org.springframework.beans.factory.annotation.Value
import org.springframework.stereotype.Component
import org.vosk.Model
import org.vosk.Recognizer
import java.io.*
import java.nio.ByteBuffer
import javax.sound.sampled.*
import kotlin.experimental.and
import kotlin.jvm.Throws
import org.vosk.LibVosk
import org.vosk.LogLevel
import ws.schild.jave.Encoder
import ws.schild.jave.MultimediaObject
import ws.schild.jave.encode.AudioAttributes
import ws.schild.jave.encode.EncodingAttributes


/**
 * 语音识别为文字
 *
 * @author 祁雪
 */
@Component
open class VoiceToText {

    @Value("\${com.bot.voice.vos.path}")
    open fun setVosPath(vosPath: String) {
        VoiceToText.vosPath = vosPath
    }

    companion object {

        lateinit var vosPath: String

        fun toText(file: File): String? {

            val ais = voiceTo16KHz(file.readBytes())
//            val ais = AudioSystem.getAudioInputStream(file)
/*
            val rdf = RandomAccessFile(file, "r")
            val track = rdf.use {
                log.info("声音尺寸:{}", toInt(read(it, 4, 4)));
                log.info("音频格式:{}", toShort(read(it, 20, 2)));
                log.info("采样率、音频采样级别 16000 = 16KHz: {}", toInt(read(it, 24, 4)));
                log.info("每秒波形的数据量：{}", toShort(read(it, 22, 2)));
                log.info("采样帧的大小：{}", toShort(read(it, 32, 2)));
                log.info("采样位数：{}", toShort(read(it, 34, 2)));
                return@use toShort(read(it, 22, 2))
            }
            log.info("1 单声道 2 双声道: {}", track);*/

            LibVosk.setLogLevel(LogLevel.INFO)
            log.info("启动语音识别")

            val model = Model(vosPath)

            var resultMap: HashMap<String, Any>? = null
            model.run {
                ais.use { aiss ->
//                    val recognizer = Recognizer(this, (16000f * track))
                    val recognizer = Recognizer(this, 16000f * aiss.format.channels)

                    recognizer.run {

                        var read: Int = -1
//                        val byteArray = ByteArray(1024)
                        val byteArray = ByteArray(512)
                        while ((aiss.read(byteArray).also { read = it }) != -1) {
                            this.acceptWaveForm(byteArray, read)
                        }

                        val result = this.finalResult
                        log.info("语音识别结束 => ${result}")
                        /*resultMap = Gson().fromJson(
                            this.finalResult,
                            object : TypeToken<HashMap<String, Any>>() {}.type
                        ) as HashMap<String, Any>*/

                        resultMap = ObjectMapper().readValue(
                            result,
                            object : TypeReference<HashMap<String, Any>>() {}
                        )

                        this
                    }.close()
                }
                return@run this
            }.close()

            if (resultMap == null || !resultMap!!.containsKey("text")) {
                return null
            }

            return resultMap!!.get("text").toString()

        }

        /**
         * 将音频文件转为16khz采集率
         */
        @Throws(IOException::class, UnsupportedAudioFileException::class)
        private fun voiceTo16KHz(byteArray: ByteArray): AudioInputStream {

            val audioIn: AudioInputStream = AudioSystem.getAudioInputStream(
                ByteArrayInputStream(byteArray)
            )

            audioIn.use {
                val srcFormat: AudioFormat = it.getFormat()

                val targetSampleRate = 16000f
                val dstFormat = AudioFormat(
                    srcFormat.getEncoding(),
                    targetSampleRate,
                    srcFormat.getSampleSizeInBits(),
                    srcFormat.getChannels(),
                    srcFormat.getFrameSize(),
                    srcFormat.getFrameRate(),
                    srcFormat.isBigEndian()
                )
                val convertedIn: AudioInputStream = AudioSystem.getAudioInputStream(dstFormat, it)
                convertedIn.use {
                    return it
                }
            }
        }

        private fun toInt(b: ByteArray): Int {
            return ByteBuffer.wrap(b).int
        }

        private fun toShort(b: ByteArray): Short {
            return ByteBuffer.wrap(b).short
        }

        @Throws(IOException::class)
        private fun read(rdf: RandomAccessFile, pos: Long, length: Int): ByteArray {
            rdf.seek(pos)
            val result = ByteArray(length)
            for (i in 0 until length) {
                result[i] = rdf.readByte()
            }
            return result
        }

        fun saveMp3(sourceFile: File, targetFile: File): File? {
            log.info("开始转换音频")
            val audio = AudioAttributes()
            val encoder = Encoder()

            audio.setCodec("libmp3lame")
            audio.setSamplingRate(16000)

            val attrs = EncodingAttributes()
            attrs.setInputFormat("mp3")
            attrs.setOutputFormat("mp3")

            attrs.setAudioAttributes(audio)

            try {
                val source = MultimediaObject(sourceFile)
                encoder.encode(source, targetFile, attrs)

                log.info("音频转换完成")
                return targetFile
            } catch (e: Exception) {
                log.error("音频转换失败", e)
            }

            return null
        }

        fun saveToWav(sourceFile: File, targetFile: File): File? {

            val byteArrayInputStream = ByteArrayInputStream(sourceFile.readBytes())

            byteArrayInputStream.use {

                val sourceAIS = AudioSystem.getAudioInputStream(it)

                val sourceFormat = sourceAIS.format

                val mp3tFormat = AudioFormat(
                    AudioFormat.Encoding.PCM_SIGNED,
                    sourceFormat.getSampleRate(),
                    16,
                    sourceFormat.getChannels(),
                    sourceFormat.getChannels() * 2,
                    sourceFormat.getSampleRate(),
                    false
                )
                val mp3AIS = AudioSystem.getAudioInputStream(mp3tFormat, sourceAIS)

                val pcmFormat = AudioFormat(
                    AudioFormat.Encoding.PCM_SIGNED,
                    16000f,
                    16,
                    1,
                    2,
                    16000f,
                    false
                )

                val pcmAIS = AudioSystem.getAudioInputStream(pcmFormat, mp3AIS)

                targetFile.parentFile.mkdirs()

                AudioSystem.write(pcmAIS, AudioFileFormat.Type.WAVE, targetFile)

            }

            return if (targetFile.exists()) targetFile else null
        }

    }

}

fun main() {
    VoiceToText.vosPath = "D:\\xun lian rep\\vosk-model-small-cn-0.22"
    /*val file = VoiceToText.saveToWav(
        File("C:\\Users\\76167\\Desktop\\文档数据\\test\\a4f5aff174b3445a3a6402f1ecb2652f.mp3"),
        File("C:\\Users\\76167\\Desktop\\文档数据\\test\\1.wav")
    )
    if (file == null) {
        return
    }
    println(file.absolutePath)*/
    val file = File(
        "C:\\Users\\76167\\Desktop\\文档数据\\test\\" +
                "015476657d110160b8a68e47bfa74b85.mp3.wav"
    )
    val text = VoiceToText.toText(file)
    println("识别内容 => ${text}")
}